import asyncio
import os
import dotenv
import random

from openai import AsyncOpenAI
from openai.types.responses import ResponseTextDeltaEvent

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
    ItemHelpers,
    WebSearchTool,
)

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()
agent = Agent(
    name="Assistant",
    tools=[
        WebSearchTool()
    ],
)

async def main():
    result = await Runner.run(agent, "考虑到我的喜好和今天旧金山的天气，我应该去哪家咖啡店？", run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER))
    print(result.final_output)

if __name__ == "__main__":
    asyncio.run(main())